'''
Author: Kivvvvi 3385856680@qq.com
Date: 2024-11-08 15:14:00
LastEditors: Kivvvvi
LastEditTime: 2024-11-10 21:06:50
Description: hft_api
'''
import requests
import pandas as pd

api_url = "http://8.148.27.197:8000"

def fetch_data_matrix(sql_query):
    try:
        payload = {"sql_query": sql_query}
        response = requests.post(f"{api_url}/query", json=payload)

        if response.status_code == 200:
            result = response.json()
            if result.get("status") == "success":
                return result["data"]
            else:
                return "Error: Failed to fetch data."
        elif response.status_code == 400:
            return "Error: Invalid query. Please check your SQL syntax or query safety."
        elif response.status_code == 500:
            return "Error: Server error. Query execution failed."
        else:
            return f"Error: Unexpected status code {response.status_code}."

    except requests.exceptions.RequestException as e:
        return "Error: Failed to connect to server." + f"Details: {e}"
    
    
class Querier:
    @staticmethod
    def get_stocks_daily(columns=None, codes=None, start=None, end=None):
        sql_query = "SELECT {} FROM daily_stock_data".format(",".join(columns) if columns else "*")
        conditions = []
        if codes:
            codes_str = ", ".join(f"'{code}'" for code in codes)
            conditions.append(f"stock_code IN ({codes_str})")
        if start:
            conditions.append(f"date >= '{start}'")
        if end:
            conditions.append(f"date <= '{end}'")
        
        if conditions:
            sql_query += " WHERE " + " AND ".join(conditions)
        
        data = fetch_data_matrix(sql_query)
        if isinstance(data, str) and data.startswith("Error"):
            return data  
        df = pd.DataFrame(data, columns=columns)
        return df

    @staticmethod
    def get_stocks_pivot(value="close", codes=None, start=None, end=None):
        df = Querier.get_stocks_daily(["stock_code", "date", value], codes, start, end)
        
        if isinstance(df, str) and df.startswith("Error"):
            return df
        
        pivot_df = df.pivot_table(values=value, index="date", columns="stock_code")
        return pivot_df